import numpy as np

# 定义计算Berge稳定性的函数
def compute_berge_stabilities(EJ1, EJ1plus, EJ2, EJ2plus, EP1plus, EP2plus, EJ3=None, EJ3plus=None, EP3plus=None):
    # Number of states
    N = int(np.sqrt(len(EJ1)))

    # Diagonal matrix
    D = np.eye(N)

    # Accessibility matrices for DM1 and DM2 (and DM3 if provided)
    J1 = np.array(EJ1).reshape(N, N)
    J1plus = np.array(EJ1plus).reshape(N, N)
    J2 = np.array(EJ2).reshape(N, N)
    J2plus = np.array(EJ2plus).reshape(N, N)
    
    if EJ3 is not None:
        J3 = np.array(EJ3).reshape(N, N)
        J3plus = np.array(EJ3plus).reshape(N, N)
    else:
        J3 = None
        J3plus = None

    # Strict preference matrices for DM1 and DM2 (and DM3 if provided)
    P1plus = np.array(EP1plus).reshape(N, N)
    P2plus = np.array(EP2plus).reshape(N, N)
    
    if EP3plus is not None:
        P3plus = np.array(EP3plus).reshape(N, N)
    else:
        P3plus = None

    # Matrix of ones
    E = np.ones((N, N))

    # Helper function to compute combined move matrix
    def compute_combined_moves(J2, J3, delta):
        M2 = J2
        M3 = J3
        M23 = np.sign(M2 + M3)
        for t in range(1, delta + 1):
            MT2 = M2
            M2 = np.sign(np.dot(J2, M3))
            M3 = np.sign(np.dot(J3, MT2))
            M23 = np.sign(M23 + np.sign(M2 + M3))
        np.fill_diagonal(M23, 0)  # Ensure the coalition does not return to the initial state
        return M23

    # Theorem 2 - Berge
    delta = np.sum(J2 != 0) + np.sum(J3 != 0) if J3 is not None else np.sum(J2 != 0)
    M23 = compute_combined_moves(J2, J3 if J3 is not None else J2, delta)
    MBerge1 = np.dot(M23, P1plus.T)
    Berge1 = np.diag(MBerge1)

    # Theorem 3 - WB
    delta = np.sum(J2plus != 0) + np.sum(J3plus != 0) if J3plus is not None else np.sum(J2plus != 0)
    M23 = compute_combined_moves(J2plus, J3plus if J3plus is not None else J2plus, delta)
    MWB1 = np.dot(M23, P1plus.T)
    WB1 = np.diag(MWB1)

    # Theorem 4 - MB
    delta = np.sum(J2 != 0) + np.sum(J3 != 0) if J3 is not None else np.sum(J2 != 0)
    M23 = compute_combined_moves(J2, J3 if J3 is not None else J2, delta)
    MMB1 = np.dot((J1 + D), np.dot(M23, P1plus.T))
    MB1 = np.diag(MMB1)

    # Theorem 5 - SMB
    delta = np.sum(J2 != 0) + np.sum(J3 != 0) if J3 is not None else np.sum(J2 != 0)
    M23 = compute_combined_moves(J2, J3 if J3 is not None else J2, delta)
    MSMB1 = np.dot((J1 + D), np.dot(M23, (P1plus.T + np.dot(J1, P1plus.T))))
    SMB1 = np.diag(MSMB1)

    # Theorem 6 - WMB
    delta = np.sum(J2plus != 0) + np.sum(J3plus != 0) if J3plus is not None else np.sum(J2plus != 0)
    M23 = compute_combined_moves(J2plus, J3plus if J3plus is not None else J2plus, delta)
    MWMB1 = np.dot((J1 + D), np.dot(M23, P1plus.T))
    WMB1 = np.diag(MWMB1)

    # Theorem 7 - WSMB
    delta = np.sum(J2plus != 0) + np.sum(J3plus != 0) if J3plus is not None else np.sum(J2plus != 0)
    M23 = compute_combined_moves(J2plus, J3plus if J3plus is not None else J2plus, delta)
    MWSMB1 = np.dot((J1 + D), np.dot(M23, (P1plus.T + np.dot(J1, P1plus.T))))
    WSMB1 = np.diag(MWSMB1)

    # Aggregate results
    results = np.column_stack((Berge1, WB1, MB1, SMB1, WMB1, WSMB1))
    return results

# Example usage for bilateral conflict (2 DMs)
EJ1 = [0, 1, 0, 0, 1, 1, 0, 0, 0]  # Example data for J1
EJ1plus = [0, 1, 0, 0, 1, 1, 0, 0, 0]  # Example data for J1plus
EJ2 = [0, 0, 1, 0, 0, 0, 0, 0, 1]  # Example data for J2
EJ2plus = [0, 0, 1, 0, 0, 0, 0, 0, 1]  # Example data for J2plus
EP1plus = [0, 1, 0, 0, 0, 1, 0, 0, 0]  # Example data for P1plus
EP2plus = [0, 0, 1, 0, 0, 0, 0, 0, 1]  # Example data for P2plus

# Compute Berge stabilities for bilateral conflict
results_bilateral = compute_berge_stabilities(EJ1, EJ1plus, EJ2, EJ2plus, EP1plus, EP2plus)
print("Berge Stabilities Results for Bilateral Conflict:")
print(results_bilateral)

# Example usage for trilateral conflict (3 DMs)
EJ3 = [0, 0, 0, 1, 0, 0, 0, 0, 0]  # Example data for J3
EJ3plus = [0, 0, 0, 1, 0, 0, 0, 0, 0]  # Example data for J3plus
EP3plus = [0, 0, 0, 1, 0, 0, 0, 0, 0]  # Example data for P3plus

# Compute Berge stabilities for trilateral conflict
results_trilateral = compute_berge_stabilities(EJ1, EJ1plus, EJ2, EJ2plus, EP1plus, EP2plus, EJ3, EJ3plus, EP3plus)
print("Berge Stabilities Results for Trilateral Conflict:")
print(results_trilateral)
